{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1 align=\"center\">Focused Ion Beam Scanning Electron Microscopy Image Segmentation</h1>\n",
    "\n",
    "\n",
    "**Summary:**\n",
    "1. SimpleITK supports a large number of filters that facilitate classical segmentation algorithms (variety of thresholding algorithms, watersheds...).\n",
    "2. Once your data is segmented SimpleITK enables you to efficiently post process the segmentation (e.g. label distinct objects, analyze object shapes).\n",
    "\n",
    "This notebook will illustrate the use of SimpleITK for segmentation of bacteria from a 3D Focused Ion Beam Scanning Electron Microscopy (FIB-SEM) image. The specific bacterium is <a href=\"https://en.wikipedia.org/wiki/Bacillus_subtilis\">bacillus subtilis</a>, a rod shaped organism naturally found in soil and plants. The bacteria have been subjected to stress to initiate the process of forming an endospore. These endospores can be seen as a generally dark ellipsoid inside the individual bacterium. \n",
    "\n",
    "Data provided courtesy of the Laboratory of Cell Biology, Center for Cancer Research, National Cancer Institute, National Institutes of Health."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import SimpleITK as sitk\n",
    "import pandas as pd\n",
    "\n",
    "%matplotlib notebook\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import gui\n",
    "from math import ceil\n",
    "\n",
    "%run update_path_to_download_script\n",
    "from downloaddata import fetch_data as fdata"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load data\n",
    "\n",
    "Load the 3D volume and display it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "img = sitk.ReadImage(fdata(\"fib_sem_bacillus_subtilis.mha\"))\n",
    "gui.MultiImageDisplay(image_list = [img], figure_size=(8,4));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Segmentation\n",
    "\n",
    "To allow us to analyze the shape of whole bacteria we first need to segment them. We will do this in several steps:\n",
    "1. Separate the bacteria from the embedding resin background.\n",
    "2. Mark each potential bacterium with a unique label, to evaluate the segmentation.\n",
    "3. Remove small components and fill small holes using binary morphology operators (opening and closing).\n",
    "4. Use seed based watersheds to perform final segmentation.\n",
    "5. Remove bacterium that are connected to the image boundary."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Separate the bacteria from the background\n",
    "\n",
    "Based on the visualization of the data above, it intuitively appears that the background and foreground are separable using a single intensity threshold. Our first step towards validating this observation is to plot the intensity distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.hist(sitk.GetArrayViewFromImage(img).flatten(), bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The histogram is bi-modal with a clear separation, which we have manually identified as having an intensity value of 120.\n",
    "\n",
    "We can also use one of several binary threshold selection filters available in SimpleITK. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "threshold_filters = {'Otsu': sitk.OtsuThresholdImageFilter(),\n",
    "                     'Triangle' : sitk.TriangleThresholdImageFilter(),\n",
    "                     'Huang' : sitk.HuangThresholdImageFilter(),\n",
    "                     'MaxEntropy' : sitk.MaximumEntropyThresholdImageFilter()}\n",
    "\n",
    "filter_selection = 'Manual'\n",
    "try:\n",
    "  thresh_filter = threshold_filters[filter_selection]\n",
    "  thresh_filter.SetInsideValue(0)\n",
    "  thresh_filter.SetOutsideValue(1)\n",
    "  thresh_img = thresh_filter.Execute(img)\n",
    "  thresh_value = thresh_filter.GetThreshold()\n",
    "except KeyError:\n",
    "  thresh_value = 120\n",
    "  thresh_img = img>thresh_value\n",
    "\n",
    "print(\"Threshold used: \" + str(thresh_value))    \n",
    "gui.MultiImageDisplay(image_list = [sitk.LabelOverlay(img, thresh_img)],                   \n",
    "                      title_list = ['Binary Segmentation'], figure_size=(8,4));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mark each potential bacterium with unique label and evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "stats = sitk.LabelShapeStatisticsImageFilter()\n",
    "stats.Execute(sitk.ConnectedComponent(thresh_img))\n",
    "\n",
    "# Look at the distribution of sizes of connected components (bacteria).\n",
    "label_sizes = [ stats.GetNumberOfPixels(l) for l in stats.GetLabels() if l != 1]\n",
    "\n",
    "plt.figure()\n",
    "plt.hist(label_sizes,bins=200)\n",
    "plt.title(\"Distribution of Object Sizes\")\n",
    "plt.xlabel(\"size in pixels\")\n",
    "plt.ylabel(\"number of objects\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The histogram above shows tens of thousands of very small labels which are not visually detected by looking at the segmentation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Remove small islands and holes\n",
    "\n",
    "Using binary morphological operations we remove small objects using the opening operation and fill small holes using the closing operation. The use of opening and closing by reconstruction maintains the boundary of the original objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "cleaned_thresh_img = sitk.BinaryOpeningByReconstruction(thresh_img, [10, 10, 10])\n",
    "cleaned_thresh_img = sitk.BinaryClosingByReconstruction(cleaned_thresh_img, [10, 10, 10])\n",
    "\n",
    "gui.MultiImageDisplay(image_list = [sitk.LabelOverlay(img, cleaned_thresh_img)],                   \n",
    "                      title_list = ['Cleaned Binary Segmentation'], figure_size=(8,4));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check that the number of objects defined by the binary image is more reasonable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "stats = sitk.LabelShapeStatisticsImageFilter()\n",
    "stats.Execute(sitk.ConnectedComponent(cleaned_thresh_img))\n",
    "\n",
    "# Look at the distribution of sizes of connected components (bacteria).\n",
    "label_sizes = [ stats.GetNumberOfPixels(l) for l in stats.GetLabels() if l != 1]\n",
    "\n",
    "plt.figure()\n",
    "plt.hist(label_sizes,bins=200)\n",
    "plt.title(\"Distribution of Object Sizes\")\n",
    "plt.xlabel(\"size in pixels\")\n",
    "plt.ylabel(\"number of objects\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After the morphological operations, our binary image seems to have a reasonable number of objects, but is this true? We next look at the unique objects defined by this binary segmentation (each object is marked with a unique color)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "gui.MultiImageDisplay(image_list = [sitk.LabelOverlay(img, sitk.ConnectedComponent(cleaned_thresh_img))],                   \n",
    "                      title_list = ['Cleaned Binary Segmentation'],figure_size=(8,4));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Seed based watershed segmentation\n",
    "\n",
    "The bacteria appear to be segmented correctly from the background but not from each other. Using the visualization and histogram above we see that in 3D many of them are connected, even if on a slice by slice inspection they appear separate.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dist_img = sitk.SignedMaurerDistanceMap(cleaned_thresh_img != 0, insideIsPositive=False, squaredDistance=False, useImageSpacing=False)\n",
    "radius = 10\n",
    "# Seeds have a distance of \"radius\" or more to the object boundary, they are uniquely labelled.\n",
    "seeds = sitk.ConnectedComponent(dist_img < -radius)\n",
    "# Relabel the seed objects using consecutive object labels while removing all objects with less than 15 pixels.\n",
    "seeds = sitk.RelabelComponent(seeds, minimumObjectSize=15)\n",
    "# Run the watershed segmentation using the distance map and seeds.\n",
    "ws = sitk.MorphologicalWatershedFromMarkers(dist_img, seeds, markWatershedLine=True)\n",
    "ws = sitk.Mask( ws, sitk.Cast(cleaned_thresh_img, ws.GetPixelID()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the distance map, the unique seeds and final object segmentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "gui.MultiImageDisplay(image_list = [dist_img,\n",
    "                                    sitk.LabelOverlay(img, seeds),\n",
    "                                    sitk.LabelOverlay(img, ws)],                   \n",
    "                      title_list = ['Segmentation Distance',\n",
    "                                    'Watershed Seeds',\n",
    "                                    'Binary Watershed Labeling'],\n",
    "                      shared_slider=True,\n",
    "                      horizontal=False,\n",
    "                      figure_size=(6,12));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Removal of objects touching the image boundary\n",
    "\n",
    "We are not sure objects touching the image boundary are whole bacteria, so we remove them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# The image has a small black border which we account for here.\n",
    "bgp = sitk.BinaryGrindPeak( (ws!=0)| (img==0))\n",
    "non_border_seg = sitk.Mask( ws, bgp==0)\n",
    "gui.MultiImageDisplay(image_list = [sitk.LabelOverlay(img, non_border_seg)],                   \n",
    "                      title_list = ['Final Segmentation'],figure_size=(8,4));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Object Analysis\n",
    "\n",
    "Once we have the segmented objects we look at their shapes and the intensity distributions inside the objects.\n",
    "\n",
    "Note that sizes are in nanometers. ITK and consequently SimpleITK are agnostic of the actual measurement units. It is up to you as the developer to explicitly use the correct units and more importantly, <a href=\"https://en.wikipedia.org/wiki/Mars_Climate_Orbiter\">DO NOT MIX UNITS</a>.\n",
    "\n",
    "We first compute all of the measurements we are interested in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "shape_stats = sitk.LabelShapeStatisticsImageFilter()\n",
    "shape_stats.ComputeOrientedBoundingBoxOn()\n",
    "shape_stats.Execute(non_border_seg)\n",
    "\n",
    "intensity_stats = sitk.LabelIntensityStatisticsImageFilter()\n",
    "intensity_stats.Execute(non_border_seg,img) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Insert the values into a pandas dataframe and display some descriptive statistics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "stats_list = [ (shape_stats.GetPhysicalSize(i),\n",
    "               shape_stats.GetElongation(i),\n",
    "               shape_stats.GetFlatness(i),\n",
    "               shape_stats.GetOrientedBoundingBoxSize(i)[0],\n",
    "               shape_stats.GetOrientedBoundingBoxSize(i)[2],\n",
    "               intensity_stats.GetMean(i),\n",
    "               intensity_stats.GetStandardDeviation(i),\n",
    "               intensity_stats.GetSkewness(i)) for i in shape_stats.GetLabels()]\n",
    "cols=[\"Volume (nm^3)\",\n",
    "      \"Elongation\",\n",
    "      \"Flatness\",\n",
    "      \"Oriented Bounding Box Minimum Size(nm)\",\n",
    "      \"Oriented Bounding Box Maximum Size(nm)\",\n",
    "     \"Intensity Mean\",\n",
    "     \"Intensity Standard Deviation\",\n",
    "     \"Intensity Skewness\"]\n",
    "\n",
    "# Create the pandas data frame and display descriptive statistics.\n",
    "stats = pd.DataFrame(data=stats_list, index=shape_stats.GetLabels(), columns=cols)\n",
    "stats.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a plot to investigate the relationship, possible correlations, between volume and object shape characteristics (elongation, flatness, principal moments). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(nrows=len(cols), ncols=2, figsize=(6,4*len(cols)))\n",
    "axes[0,0].axis('off')\n",
    "\n",
    "stats.loc[:,cols[0]].plot.hist(ax=axes[0,1], bins=25)\n",
    "axes[0,1].set_xlabel(cols[0])\n",
    "axes[0,1].xaxis.set_label_position(\"top\")\n",
    "\n",
    "for i in range(1,len(cols)):\n",
    "    c = cols[i]\n",
    "    bar = stats.loc[:,[c]].plot.hist(ax=axes[i,0], bins=20,orientation='horizontal',legend=False)\n",
    "    bar.set_ylabel(stats.loc[:,[c]].columns.values[0])    \n",
    "    scatter = stats.plot.scatter(ax=axes[i,1],y=c,x=cols[0])\n",
    "    scatter.set_ylabel('')\n",
    "    # Remove axis labels from all plots except the last (they all share the labels)\n",
    "    if(i<len(cols)-1):\n",
    "        bar.set_xlabel('')\n",
    "        scatter.set_xlabel('')\n",
    "# Adjust the spacing between plot columns and set the plots to have a tight\n",
    "# layout inside the figure.\n",
    "plt.subplots_adjust(wspace=0.4)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we visualize a lineup of the bacteria using a coordinate system that is defined by the oriented bounding box enclosing each of them. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "bacteria_labels = shape_stats.GetLabels()\n",
    "bacteria_volumes = [shape_stats.GetPhysicalSize(label) for label in bacteria_labels] \n",
    "num_images = 5 # number of bacteria images we want to display\n",
    "\n",
    "bacteria_labels_volume_sorted = [label for _,label in sorted(zip(bacteria_volumes, bacteria_labels))]\n",
    "\n",
    "resampler = sitk.ResampleImageFilter()\n",
    "aligned_image_spacing = [10,10,10] #in nanometers\n",
    "\n",
    "for label in bacteria_labels_volume_sorted[0:num_images]:\n",
    "    aligned_image_size = [ int(ceil(shape_stats.GetOrientedBoundingBoxSize(label)[i]/aligned_image_spacing[i])) for i in range(3) ]\n",
    "    direction_mat = shape_stats.GetOrientedBoundingBoxDirection(label)\n",
    "    aligned_image_direction = [direction_mat[0], direction_mat[3], direction_mat[6], \n",
    "                               direction_mat[1], direction_mat[4], direction_mat[7],\n",
    "                               direction_mat[2], direction_mat[5], direction_mat[8] ] \n",
    "    resampler.SetOutputDirection(aligned_image_direction)\n",
    "    resampler.SetOutputOrigin(shape_stats.GetOrientedBoundingBoxOrigin(label))\n",
    "    resampler.SetOutputSpacing(aligned_image_spacing)\n",
    "    resampler.SetSize(aligned_image_size)\n",
    "    \n",
    "    obb_img = resampler.Execute(img)\n",
    "    # Change the image axes order so that we have a nice display.\n",
    "    obb_img = sitk.PermuteAxes(obb_img,[2,1,0])\n",
    "    gui.MultiImageDisplay(image_list = [obb_img],                   \n",
    "                          title_list = [\"OBB_{0}\".format(label)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
